import os
import numpy as np
import xml.etree.ElementTree as ET

try:
    from .cvio import cvio
except:
    from cvio import cvio


def indent(elem, level=0):
    i = "\n" + level*"\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            indent(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i


def generateXML(outPath, filename, w, h, bboxes, labels):

    root = ET.Element('annotation')
    xfolder = ET.SubElement(root, 'folder')
    xfolder.text = 'images'
    xfilename = ET.SubElement(root, 'filename')
    xfilename.text = filename
    xpath = ET.SubElement(root, 'path')
    xpath.text = 'images/'+filename
    xsource = ET.SubElement(root, 'source')
    xdatabase = ET.SubElement(xsource, 'database')
    xdatabase.text = 'Unknown'

    xsize = ET.SubElement(root, 'size')
    xwidth = ET.SubElement(xsize, 'width')
    xwidth.text = str(w)
    xheight = ET.SubElement(xsize, 'height')
    xheight.text = str(h)
    depth = ET.SubElement(xsize, 'depth')
    depth.text = '3'

    xsegmented = ET.SubElement(root, 'segmented')
    xsegmented.text = '0'

    for i in range(len(bboxes)):
        xxobject = ET.SubElement(root, 'object')
        xxname = ET.SubElement(xxobject, 'name')
        xxname.text = labels[i]
        xxpose = ET.SubElement(xxobject, 'pose')
        xxpose.text = 'Unspecified'
        xxdifficult = ET.SubElement(xxobject, 'difficult')
        xxdifficult.text = '0'
        xxtruncated = ET.SubElement(xxobject, 'truncated')
        xxtruncated.text = '0'
        xxbndbox = ET.SubElement(xxobject, 'bndbox')
        xxmin = ET.SubElement(xxbndbox, 'xmin')
        xxmin.text = str(int(bboxes[i][0]))
        xymin = ET.SubElement(xxbndbox, 'ymin')
        xymin.text = str(int(bboxes[i][1]))
        xxmax = ET.SubElement(xxbndbox, 'xmax')
        xxmax.text = str(int(bboxes[i][2]))
        xymax = ET.SubElement(xxbndbox, 'ymax')
        xymax.text = str(int(bboxes[i][3]))

    tree = ET.ElementTree(root)
    indent(root)

    tree.write(outPath, encoding='utf-8')


def ply2box(ply):
    ply = np.array(ply)
    xmin = ply[:, 0].min()
    xmax = ply[:, 0].max()
    ymin = ply[:, 1].min()
    ymax = ply[:, 1].max()
    return [xmin, ymin, xmax, ymax]


def get_bboxes(annInfo):
    shapes = annInfo['shapes']
    bboxes = []
    labels = []
    for shape in shapes:
        points = shape['points']
        label = shape['label']
        bboxes.append(ply2box(points))
        labels.append(label)

    return bboxes, labels


def mask2bbox(src, dst=''):
    ann_list = cvio.load_ext_list(src, ext_type='.json')
    n = len(ann_list)
    if n < 1:
        print('无标注数据(json).')
        return

    for i, ann in enumerate(ann_list, 1):
        f = os.path.basename(ann)
        print('[%d/%d] %s' % (i, n, f))
        annInfo=cvio.load_ann(ann)
        bboxes, labels=get_bboxes(annInfo)
        filename=os.path.basename(annInfo['imagePath'])
        w=annInfo['imageWidth']
        h=annInfo['imageHeight']
        if dst in (None, ''):
            outPath=os.path.splitext(ann)[0] + '.xml'
        else:
            outPath=os.path.join(dst, os.path.splitext(f)[0] + '.xml')
        generateXML(outPath, filename, w, h, bboxes, labels)
